from models.model import Model
from torch import nn
import torch.nn.functional as F


class CNN4(Model):
    is_relu:bool = True
    fc1_in_features:int

    def __init__(self, dst: str, is_relu:bool = True):
        super(CNN4,self).__init__(dst)
        self.is_relu = is_relu
        self.conv1 = nn.Conv2d(self.in_channels, 12, kernel_size=5, padding=2, stride=2)
        self.conv2 = nn.Conv2d(12, 12, kernel_size=5, padding=2, stride=2)
        self.conv3 = nn.Conv2d(12, 12, kernel_size=5, padding=2, stride=1)
        self.conv4 = nn.Conv2d(12, 12, kernel_size=5, padding=2, stride=1)
        self.fc = nn.Linear(self.fc1_in_features,10)


    def forward(self,x):
        out = None
        if self.is_relu:
            out = F.relu(self.conv1(x))
            out = F.relu(self.conv2(out))
            out = F.relu(self.conv3(out))
            out = F.relu(self.conv4(out))
        else:
            out = F.sigmoid(self.conv1(x))
            out = F.sigmoid(self.conv2(out))
            out = F.sigmoid(self.conv3(out))
            out = F.sigmoid(self.conv4(out))
        out = out.view(out.size(0),-1)
        out = self.fc(out)

        return out


    def _get_fc1_in_features_size(self):
        c1:int = int((self.data_size - 1)/2) + 1#c1 = c2 = c3
        c4:int = int((c1 - 1)/2) + 1

        self.fc1_in_features = c4 * c4 * 12


class CNN2(Model):
    is_relu:bool = True
    fc1_in_features:int

    def __init__(self, dst: str, is_relu:bool = True):
        super(CNN2,self).__init__(dst)
        self.conv1 = nn.Conv2d(self.in_channels, 10, kernel_size=5)
        self.pool1 = nn.MaxPool2d(2,2)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.pool2 = nn.MaxPool2d(2,2)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(self.fc1_in_features, 50)
        self.fc2 = nn.Linear(50, 10)


    def forward(self,x):
        out = None
        if self.is_relu:
            out = F.relu(self.conv1(x))
            out = self.pool1(out)
            out = F.relu(self.conv2_drop(self.conv2(out)))
            out = self.pool2(out)
        else:
            out = F.sigmoid(self.conv1(x))
            out = self.pool1(out)
            out = F.sigmoid(self.conv2_drop(self.conv2(out)))
            out = self.pool2(out)
        out = out.view(out.size(0),-1)
        out = F.relu(self.fc1(out))
        out = self.fc2(out)
        return out


    def _get_fc1_in_features_size(self):
        c1:int = int(self.data_size/2) - 2
        c2:int = int(c1/2) - 2

        self.fc1_in_features = c2 * c2 * 20
